from typing import List, Dict, Any, Optional
from datetime import datetime, timedelta
import threading
import time
from utils.logger import Logger
from utils.exceptions import ToolError
from middleware.redis_adapter import get_redis_adapter
from states.agent_state import AgentState, AgentType, AgentHealthStatus
from utils.constants import AGENT_STATUSES

# 初始化日志
logger = Logger.get_logger(agent_name="health_check_agent")


class HealthCheckAgent:
    """健康检查Agent：定期监控核心Agent状态，更新到Redis"""

    def __init__(self):
        self.redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
        self.agent_prefix = "agent:state:"  # Agent状态存储键前缀
        self.check_interval = 10  # 健康检查间隔（10秒，与Agent心跳一致）
        self.timeout = 3  # 每个Agent健康检查超时时间（3秒）
        self.max_offline_time = 30  # Agent离线阈值（30秒无响应则标记为DOWN）
        # 使用标准AgentType枚举值，与AgentState类保持一致
        self.agent_types = [agent_type.value for agent_type in AgentType] if AgentType else ["text_qa", "multi_modal", "tool_call"]
        # 健康检查线程（后台运行）
        self.check_thread = threading.Thread(target=self._run_health_check, daemon=True)
        self.running = False
        logger.info("健康检查Agent初始化成功，检查间隔: 10秒")

    def start(self) -> None:
        """启动健康检查线程"""
        if not self.running:
            self.running = True
            self.check_thread.start()
            logger.info("健康检查Agent启动成功（后台线程）")
        else:
            logger.warning("健康检查Agent已在运行中")

    def stop(self) -> None:
        """停止健康检查线程"""
        self.running = False
        self.check_thread.join(timeout=5)
        logger.info("健康检查Agent已停止")

    def _run_health_check(self) -> None:
        """健康检查循环（后台线程）"""
        while self.running:
            try:
                logger.debug("开始新一轮健康检查")
                # 遍历所有Agent类型，检查该类型下的所有实例
                for agent_type in self.agent_types:
                    self._check_agent_type(agent_type)
                # 清理长时间离线的Agent实例
                self._clean_offline_agents()
                logger.debug("本轮健康检查完成")
            except Exception as e:
                logger.error(f"健康检查循环异常: error={str(e)}", exc_info=True)
            # 等待下一轮检查
            time.sleep(self.check_interval)

    def _check_agent_type(self, agent_type: str) -> None:
        """检查指定类型的所有Agent实例"""
        # 使用标准的AgentState格式查找该类型下的所有实例：agent:state:{agent_type}:*
        pattern = f"{self.agent_prefix}{agent_type}:*"
        instance_keys = self.redis_client.get_client().keys(pattern)
        
        if not instance_keys:
            logger.debug(f"无{agent_type}类型的Agent实例")
            return

        for key in instance_keys:
            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
            # 从键名中提取实例ID：agent:state:{agent_type}:{instance_id}
            parts = key_str.split(":")
            if len(parts) < 4:
                logger.warning(f"无效的Agent状态键: {key_str}")
                continue
            
            instance_id = parts[3]
            
            # 获取Agent实例信息
            agent_info = self.redis_client.hgetall(key_str)
            if not agent_info:
                logger.warning(f"Agent实例信息不存在，跳过检查: instance_id={instance_id}")
                continue

            agent_info = self._parse_agent_info(agent_info)
            # 构建服务URL
            host_ip = agent_info.get("host_ip", "localhost")
            port = agent_info.get("port", 8000)
            service_url = f"http://{host_ip}:{port}"

            # 执行健康检查
            self._check_single_agent(instance_id, agent_type, service_url, key_str)

    def _check_single_agent(self, agent_id: str, agent_type: str, service_url: str, key_str: str) -> None:
        """检查单个Agent实例的健康状态"""
        try:
            # 构建健康检查URL
            health_url = f"{service_url}/health"
            # 发送HTTP请求（超时控制）
            import requests
            response = requests.get(health_url, timeout=self.timeout)
            response.raise_for_status()  # 抛出HTTP错误（4xx/5xx）

            # 解析健康状态响应
            health_data = response.json()
            current_time = datetime.now()
            # 更新Agent状态到Redis（使用标准AgentHealthStatus）
            self.redis_client.hset(
                key=key_str,
                mapping={
                    "status": AgentHealthStatus.HEALTHY.value,
                    "last_heartbeat": int(current_time.timestamp()),  # 使用时间戳，与AgentState一致
                    "cpu_usage": health_data.get("cpu_usage", 0.0),
                    "mem_usage": health_data.get("mem_usage", 0.0),
                    "current_tasks": health_data.get("task_count", 0)
                }
            )
            logger.debug(f"Agent健康检查通过: agent_id={agent_id}, url={health_url}, status=HEALTHY")
        except requests.exceptions.Timeout:
            # 请求超时：检查上次健康时间，判断是否标记为不健康
            self._handle_agent_timeout(key_str, agent_id)
        except Exception as e:
            # 其他异常（如连接失败、HTTP错误）
            logger.warning(f"Agent健康检查失败: agent_id={agent_id}, url={health_url}, error={str(e)}")
            self._handle_agent_timeout(key_str, agent_id)

    def _handle_agent_timeout(self, key_str: str, agent_id: str) -> None:
        """处理Agent健康检查超时"""
        last_heartbeat = self.redis_client.hget(key_str, "last_heartbeat")
        current_time = datetime.now()

        if not last_heartbeat:
            # 从未收到心跳，标记为不健康
            self.redis_client.hset(key_str, "status", AgentHealthStatus.UNHEALTHY.value)
            logger.warning(f"Agent首次检查超时，标记为不健康: agent_id={agent_id}")
            return

        # 计算上次心跳到现在的时间差（使用时间戳）
        try:
            last_heartbeat_ts = int(last_heartbeat.decode("utf-8"))
            current_ts = int(current_time.timestamp())
            offline_duration = current_ts - last_heartbeat_ts

            if offline_duration > self.max_offline_time:
                # 超过离线阈值，标记为不健康
                self.redis_client.hset(key_str, "status", AgentHealthStatus.UNHEALTHY.value)
                logger.warning(f"Agent长时间离线，标记为不健康: agent_id={agent_id}, 离线时长={offline_duration:.1f}秒")
            else:
                # 未超过阈值，保持原状态（可能是临时网络波动）
                current_status = self.redis_client.hget(key_str, "status")
                current_status = current_status.decode("utf-8") if current_status else AgentHealthStatus.UNHEALTHY.value
                logger.warning(
                    f"Agent检查超时，当前状态: {current_status}, 离线时长={offline_duration:.1f}秒: agent_id={agent_id}")
        except Exception as e:
            logger.error(f"处理Agent超时异常: agent_id={agent_id}, error={str(e)}")

    def _parse_agent_info(self, raw_info: Dict[bytes, bytes]) -> Dict[str, Any]:
        """解析Redis存储的Agent信息"""
        parsed_info = {}
        for k, v in raw_info.items():
            key = k  # 在Python 3中，键已经是字符串类型
            value = v  # 在Python 3中，值也已经是字符串类型
            # 数值类型转换
            if key in ["cpu_usage", "mem_usage"]:
                try:
                    parsed_info[key] = float(value)
                except:
                    parsed_info[key] = 0.0
            elif key in ["task_count"]:
                try:
                    parsed_info[key] = int(value)
                except:
                    parsed_info[key] = 0
            else:
                parsed_info[key] = value
        return parsed_info

    def _clean_offline_agents(self) -> None:
        """清理长时间离线的Agent实例"""
        current_time = datetime.now()
        cleaned_count = 0
        current_ts = int(current_time.timestamp())

        for agent_type in self.agent_types:
            # 使用标准格式查找该类型的所有实例
            pattern = f"{self.agent_prefix}{agent_type}:*"
            instance_keys = self.redis_client.get_client().keys(pattern)
            
            for key in instance_keys:
                key_str = key.decode("utf-8") if isinstance(key, bytes) else key
                parts = key_str.split(":")
                if len(parts) < 4:
                    continue
                
                instance_id = parts[3]
                
                # 获取最后心跳时间
                last_heartbeat = self.redis_client.hget(key_str, "last_heartbeat")
                if not last_heartbeat:
                    # 从未收到心跳，清理实例
                    self.redis_client.delete(key_str)
                    cleaned_count += 1
                    logger.debug(f"清理无心跳记录的Agent实例: {key_str}")
                    continue

                try:
                    last_heartbeat_ts = int(last_heartbeat.decode("utf-8"))
                    offline_duration = current_ts - last_heartbeat_ts
                    
                    # 超过10分钟离线，清理实例
                    if offline_duration > 600:
                        self.redis_client.delete(key_str)
                        cleaned_count += 1
                        logger.debug(f"清理长时间离线Agent实例: {key_str}, 离线时长={offline_duration:.1f}秒")
                except Exception as e:
                    logger.error(f"清理Agent实例异常: key={key_str}, error={str(e)}")

        if cleaned_count > 0:
            logger.info(f"清理长时间离线Agent实例数: {cleaned_count}")

    def _remove_agent_instance(self, agent_id: str, agent_type: str) -> None:
        """移除Agent实例"""
        # 使用标准格式构建键名
        agent_state_key = f"{self.agent_prefix}{agent_type}:{agent_id}"
        # 直接删除状态存储
        deleted = self.redis_client.delete(agent_state_key)
        if deleted > 0:
            logger.info(f"清理离线Agent实例: agent_id={agent_id}, agent_type={agent_type}")
        else:
            logger.debug(f"Agent实例不存在，跳过删除: agent_id={agent_id}, agent_type={agent_type}")

    def get_healthy_agents(self, agent_type: str) -> List[Dict[str, Any]]:
        """获取指定类型的健康Agent实例（状态为HEALTHY）"""
        healthy_agents = []
        current_ts = int(datetime.now().timestamp())
        
        # 使用标准格式查找该类型的所有实例
        pattern = f"{self.agent_prefix}{agent_type}:*"
        instance_keys = self.redis_client.get_client().keys(pattern)
        
        if not instance_keys:
            logger.debug(f"无{agent_type}类型的Agent实例")
            return []

        for key in instance_keys:
            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
            # 获取Agent实例信息
            agent_info = self.redis_client.hgetall(key_str)
            if not agent_info:
                continue

            agent_info = self._parse_agent_info(agent_info)
            # 检查状态和心跳时间
            if (agent_info.get("status") == AgentHealthStatus.HEALTHY.value and 
                agent_info.get("last_heartbeat")):
                
                # 验证心跳是否在30秒内
                try:
                    last_heartbeat_ts = int(agent_info["last_heartbeat"])
                    if current_ts - last_heartbeat_ts <= 30:
                        # 添加实例ID
                        parts = key_str.split(":")
                        if len(parts) >= 4:
                            agent_info["instance_id"] = parts[3]
                        healthy_agents.append(agent_info)
                except Exception:
                    logger.warning(f"心跳时间格式错误: {key_str}")

        # 按任务数升序排序（优先选择负载低的实例）
        healthy_agents.sort(key=lambda x: x.get("current_tasks", 0))
        logger.debug(f"获取{agent_type}健康实例数: {len(healthy_agents)}")
        return healthy_agents


# 单例实例
health_check_agent = HealthCheckAgent()

# 测试入口（本地运行时启动健康检查）
if __name__ == "__main__":
    test_logger = Logger.update_context(task_id="health_check_test", agent_name="health_check_agent")
    test_logger.info("开始测试健康检查Agent")

    # 启动健康检查线程
    health_check_agent.start()

    # 定期打印健康实例信息用于测试
    import time
    
    try:
        while True:
            time.sleep(15)
            for agent_type in [AgentType.TEXT_QA, AgentType.MULTI_MODAL, AgentType.TOOL_CALL]:
                healthy_agents = health_check_agent.get_healthy_agents(agent_type.value)
                test_logger.info(f"当前健康的{agent_type.name} Agent数量: {len(healthy_agents)}")
                for agent in healthy_agents:
                    test_logger.info(f"  - {agent['host_ip']}:{agent['port']} (任务数: {agent.get('current_tasks', 0)})")
    except KeyboardInterrupt:
        test_logger.info("测试中断，停止健康检查Agent")
    finally:
        # 停止健康检查线程
        health_check_agent.stop()
        test_logger.info("健康检查Agent测试完成")